from __future__ import absolute_import, division, print_function, unicode_literals

import datetime  # For datetime objects
import os.path  # To manage paths
import sys  # To find out the script name (in argv[0])

# Import the backtrader platform
import backtrader as bt

import MySQLDataFeed as mdf
import SimpleMySqlClass as smc
import talib as talib
import numpy as np


class MyMaMacdRis(bt.Strategy):
    params = (("period_sma10", 10), ("period_sma20", 20), ("period_sma60", 60))

    # 打印日志
    def log(self, txt, dt=None):
        dt = dt or self.data.datetime.date(0)
        print("%s, %s" % (dt, txt))

    def __init__(self):
        # 用于保存订单
        self.order = None
        # 订单价格
        self.buyprice = None
        # 订单佣金
        self.buycomm = None

        # 定义变量保存所有收盘价
        self.dataopen = self.data.open
        self.datahigh = self.data.high
        self.datalow = self.data.low
        self.dataclose = self.data.close
        self.datavolume = self.data.volume

        # 计算10日均线
        self.sma10 = bt.talib.WMA(self.dataclose, timeperiod=self.params.period_sma10)
        # 计算20日均线
        self.sma20 = bt.talib.WMA(self.dataclose, timeperiod=self.params.period_sma20)
        print("sma20 = ", self.sma20)
        # # 计算60日均线
        self.sma60 = bt.talib.WMA(self.dataclose, timeperiod=self.params.period_sma60)

        # Chaikin A/D Line 累积/派发线（Accumulation/Distribution Line） 简介：Marc Chaikin提出的一种平衡交易量指标
        self.dataad = bt.talib.AD(
            self.datahigh, self.datalow, self.dataclose, self.datavolume
        )

        # 平均趋向指数 简介：使用ADX指标，指标判断盘整、振荡和单边趋势。
        # self.dataadx = bt.talib.ADX(
        #     self.datahigh, self.datalow, self.dataclose, timeperiod=14
        # )

        # 投资学中意思为续航，指股票(或经济指数)持续增长的能力
        # self.datamom = bt.talib.MOM(self.dataclose)

        self.datamacd = bt.indicators.MACDHisto(self.datas[0])

        # 相对强弱指数
        self.dataris = bt.indicators.RSI(self.datas[0])
        print("data ris = ", self.dataris)
        # # 顺势指标
        # self.datacci = bt.talib.CCI(self.datahigh, self.datalow, self.dataclose)

        # # 资金流量指标
        # self.datamfi = bt.talib.MFI(
        #     self.datahigh, self.datalow, self.dataclose, self.datavolume
        # )

        # # 上升动向值
        # self.datamom = bt.talib.MOM(self.dataclose, timeperiod=10)

        # 定义全局变量
        self.count = 0
        for data in self.datas:
            # 转为tabib要求的数据格式
            opens = np.array(data.open.array)
            highs = np.array(data.high.array)
            lows = np.array(data.low.array)
            closes = np.array(data.close.array)
            # print(opens)
            # 数据放入self中

            # 计算十字星数据，结果为-100底部十字星，结果为100顶部十字星，0非十字星
            # res = talib.CDLDOJISTAR(opens, highs, lows, closes)
            # print("十字星数据", res)
            # self.star[data._id] = res

    def notify_order(self, order):
        # 等待订单提交、订单被cerebro接受
        if order.status in [order.Submitted, order.Accepted]:
            return

        # 等待订单完成
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(
                    "买 EXECUTED, 开仓价格: %.2f, 花费: %.2f, Comm %.2f"
                    % (order.executed.price, order.executed.value, order.executed.comm)
                )

            else:
                self.log(
                    "卖 EXECUTED, 卖出价格: %.2f, 收回资金: %.2f, Comm %.2f"
                    % (order.executed.price, order.executed.value, order.executed.comm)
                )
                self.order = None

        # 如果订单保证金不足，将不会完成，而是执行以下拒绝程序
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log("拒绝执行，Order Canceled/Margin/Rejected")

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        self.log(
            "开始交易 PROFIT, GROSS %.2f, NET %.2f" % (trade.pnl, trade.pnlcomm)
        )  # pnl：盈利  pnlcomm：手续费

    # 策略逻辑实现
    def next(self):

        # 10日均线小于20日
        cur_sma_10_lt_20 = self.sma10[0] < self.sma20[0]
        # 20日均线大于60日
        cur_sma_20_gt_60 = self.sma20[0] > self.sma60[0]

        # macd 金叉，macd零轴之上，macd上升趋势
        cur_macd_jincha = self.datamacd.histo[0] > 0
        cur_macd_gt_0 = self.datamacd.macd[0] >= 0
        cur_macd_up = self.datamacd.macd[0] > self.datamacd.macd[-1]

        # rsi 大于50，rsi上升趋势
        cur_rsi_gt_50 = self.dataris[0] > 50
        # K线当日收阳线，大于昨日上升趋势，
        cur_k_line_up = self.dataclose > self.dataopen

        cur_k_qushi = (
            self.dataclose[0] > self.dataclose[-1]
            and self.datahigh[0] > self.datahigh[-1]
            and self.datalow[0] > self.datalow[-1]
        )

        cur_jiang = self.dataclose[0] < self.datalow[-1]

        # 当今天的20日均线大于60日均线并且昨天的20日均线小于60日均线，则进入市场（买）
        # if self.sma20[0] > self.sma60[0] and self.sma20[-1] < self.sma60[-1]:
        if (
            cur_sma_20_gt_60
            and cur_macd_jincha
            and cur_macd_gt_0
            and cur_macd_up
            and cur_rsi_gt_50
            and cur_k_line_up
            and cur_k_qushi
        ):
            print("可以考虑买入，self.order = ", self.order)
            # 判断订单是否完成，完成则为None，否则为订单信息
            if not self.order is None:
                self.log("已经存在订单了")
                return

            # 若上一个订单处理完成，可继续执行买入操作
            self.order = self.buy()
            print("已经买入，self.order = ", self.order)

        # 当今天的20日均线小于60日均线并且昨天的20日均线大于60日均线，则退出市场（卖）
        elif (
            cur_sma_10_lt_20
            or not cur_sma_20_gt_60
            or not cur_macd_jincha
            or not cur_macd_gt_0
            or not cur_rsi_gt_50
            or cur_jiang
        ):

            print("需要卖出，self.order = ", self.order)
            if self.order is None:
                self.log("没有订单需要卖出")
                return
            # 卖出
            self.order = self.sell()
            print("已经卖出，self.order = ", self.order)

        for data in self.datas:
            pos = self.getposition(data).size
            # 函数出现100就代表

        # 自增处理
        self.count = self.count + 1


if __name__ == "__main__":
    # Create a cerebro entity
    cerebro = bt.Cerebro()

    # Add a strategy
    cerebro.addstrategy(MyMaMacdRis)

    tableName = "t_stock_k_data"
    sql = f"select sd_date, sd_open, sd_high, sd_low, sd_close, sd_volume from {tableName}"
    # sql = f"select sd_date as date, sd_code as code, sd_open as code, sd_high as high, sd_low as low, sd_close as low, sd_volume as volume, sd_amount as amount from {tableName}"

    data = mdf.MySQLDataFeed(sql)

    # Add the Data Feed to Cerebro
    cerebro.adddata(data, name="2Line")

    # Set our desired cash start
    cerebro.broker.setcash(1000.0)

    # Add a FixedSize sizer according to the stake
    cerebro.addsizer(bt.sizers.FixedSize, stake=10)

    # Set the commission
    cerebro.broker.setcommission(commission=0.0005)

    # Print out the starting conditions
    print("Starting Portfolio Value: %.2f" % cerebro.broker.getvalue())
    cerebro.broker.set_coc(True)  # 设置以当日收盘价成交
    # Run over everything
    cerebro.run()

    # Print out the final result
    print("Final Portfolio Value: %.2f" % cerebro.broker.getvalue())

    # Plot the result
    cerebro.plot(style="candle", barup="red", bardown="green")
